"""
RAG系统使用示例
展示如何使用文档加载器和RAG工具
"""

import sys
import os
import logging

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from rag.document_loader import DocumentLoader
from tools.rag_tool import RAGTool
from utils.config import Config

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def setup_knowledge_base():
    """设置知识库"""
    logger.info("开始设置知识库...")
    
    # 初始化文档加载器
    loader = DocumentLoader()
    
    # 处理文档
    data_dir = "/Users/pengxiong/workspace/pyWorkspace/traval-agent/data"
    collection_name = "travel_documents"
    
    success = loader.process_documents(data_dir, collection_name)
    
    if success:
        logger.info("知识库设置完成")
        
        # 显示集合信息
        info = loader.get_collection_info(collection_name)
        logger.info(f"集合信息: {info}")
        
        return True
    else:
        logger.error("知识库设置失败")
        return False

def demo_rag_usage():
    """演示RAG工具的使用"""
    logger.info("开始演示RAG工具使用...")
    
    # 初始化RAG工具
    rag_tool = RAGTool()
    
    # 示例问题
    questions = [
        "成都有哪些著名的旅游景点？",
        "成都的特色美食有哪些？",
        "去成都旅游有什么交通建议？",
        "成都的住宿情况如何？",
        "成都的气候怎么样？"
    ]
    
    logger.info("=" * 50)
    logger.info("RAG问答演示")
    logger.info("=" * 50)
    
    for i, question in enumerate(questions, 1):
        logger.info(f"\n问题 {i}: {question}")
        logger.info("-" * 30)
        
        # 获取答案
        result = rag_tool.answer_question(question)
        
        print(f"答案: {result['answer']}")
        print(f"参考文档数量: {len(result['source_documents'])}")
        
        # 显示源文档信息
        if result['source_documents']:
            print("参考文档:")
            for j, doc in enumerate(result['source_documents'][:2], 1):
                print(f"  {j}. {doc['source']}")
                print(f"     {doc['content'][:80]}...")
        
        print("\n" + "=" * 50)

def demo_search_functionality():
    """演示搜索功能"""
    logger.info("\n开始演示搜索功能...")
    
    rag_tool = RAGTool()
    
    # 搜索查询
    search_queries = [
        "成都旅游景点",
        "美食推荐",
        "交通方式"
    ]
    
    logger.info("=" * 50)
    logger.info("文档搜索演示")
    logger.info("=" * 50)
    
    for query in search_queries:
        logger.info(f"\n搜索查询: {query}")
        logger.info("-" * 30)
        
        results = rag_tool.search_documents(query, k=3)
        
        if results:
            print(f"找到 {len(results)} 个相关文档:")
            for i, result in enumerate(results, 1):
                print(f"  {i}. 相似度: {result['score']:.4f}")
                print(f"     来源: {result['metadata'].get('source', '未知')}")
                print(f"     内容: {result['content'][:80]}...")
        else:
            print("未找到相关文档")
        
        print("\n" + "=" * 50)

def main():
    """主函数"""
    logger.info("RAG系统使用示例")
    logger.info("=" * 50)
    
    try:
        # 验证配置
        Config.validate_config()
        logger.info("✓ 配置验证通过")
        
        # 设置知识库
        if not setup_knowledge_base():
            logger.error("知识库设置失败，退出示例")
            return
        
        # 演示RAG问答功能
        demo_rag_usage()
        
        # 演示搜索功能
        demo_search_functionality()
        
        logger.info("\n示例演示完成！")
        
    except Exception as e:
        logger.error(f"示例运行出错: {str(e)}")

if __name__ == "__main__":
    main()
